package com.xl.bigdata.ai.tc.test;

import com.hankcs.hanlp.classification.classifiers.IClassifier;
import com.hankcs.hanlp.classification.classifiers.NaiveBayesClassifier;
import com.hankcs.hanlp.classification.models.NaiveBayesModel;
import com.hankcs.hanlp.corpus.io.IOUtil;
import com.xl.bigdata.ai.sa.test.TestUtility;
import java.io.File;
import java.io.IOException;
import java.io.PrintStream;

public class TextClassification
{
    public static final String CORPUS_FOLDER = TestUtility.ensureTestData("搜狗文本分类语料库迷你版", "http://file.hankcs.com/corpus/sogou-text-classification-corpus-mini.zip");
    public static final String MODEL_PATH = "E://lexin/sourceCode/BoyBD/data/test/textclassification-model.ser";

    public static void main(String[] args)
            throws IOException
    {
        IClassifier classifier = new NaiveBayesClassifier(trainOrLoadModel());
        predict(classifier, "C罗获2018环球足球奖最佳球员 德尚荣膺最佳教练");
        predict(classifier, "英国造航母耗时8年仍未服役 被中国速度远远甩在身后");
        predict(classifier, "研究生考录模式亟待进一步专业化");
        predict(classifier, "如果真想用食物解压,建议可以食用燕麦");
        predict(classifier, "通用及其部分竞争对手目前正在考虑解决库存问题");
    }

    private static void predict(IClassifier classifier, String text)
    {
        System.out.printf("《%s》 属于分类 【%s】\n", new Object[] { text, classifier.classify(text) });
    }

    private static NaiveBayesModel trainOrLoadModel() throws IOException
    {
        NaiveBayesModel model = (NaiveBayesModel)IOUtil.readObjectFrom("E://lexin/sourceCode/BoyBD/data/test/textclassification-model.ser");
        if (model != null) {
            return model;
        }

        File corpusFolder = new File(CORPUS_FOLDER);
        if ((!corpusFolder.exists()) || (!corpusFolder.isDirectory()))
        {
            System.err.println("没有文本分类语料，请阅读IClassifier.train(java.lang.String)中定义的语料格式与语料下载：https://github.com/hankcs/HanLP/wiki/%E6%96%87%E6%9C%AC%E5%88%86%E7%B1%BB%E4%B8%8E%E6%83%85%E6%84%9F%E5%88%86%E6%9E%90");

            System.exit(1);
        }

        IClassifier classifier = new NaiveBayesClassifier();
        classifier.train(CORPUS_FOLDER);
        model = (NaiveBayesModel)classifier.getModel();
        IOUtil.saveObjectTo(model, "E://lexin/sourceCode/BoyBD/data/test/textclassification-model.ser");
        return model;
    }
}